import { create } from 'zustand';
import React from 'react';

interface SocketMessage {
  channel: string;
  type: string;
  clientId: string;
  body: string;
  [key: string]: unknown;
}

interface MessageHandler {
  (data: SocketMessage): void;
}

interface SharedStore {
  // 聊天设置
  chatListStyle: string;
  chatStream: boolean;
  setChatListStyle: (style: string) => void;
  setChatStream: (stream: boolean) => void;
  
  // 登录相关
  showLoginDialog: boolean;
  setShowLoginDialog: (show: boolean) => void;
  
  // WebSocket 相关
  socket: {
    conn: WebSocket | null;
    isConnected: boolean;
    messageHandlers: Map<string, MessageHandler[]>;
  };
  connectWebSocket: () => void;
  disconnectWebSocket: () => void;
  addMessageHandler: (channel: string, handler: MessageHandler) => void;
  removeMessageHandler: (channel: string) => void;
}

// 模拟 WebSocket 连接
class MockWebSocket {
  readyState = WebSocket.OPEN;
  onmessage: ((event: MessageEvent) => void) | null = null;
  onopen: ((event: Event) => void) | null = null;
  onclose: ((event: CloseEvent) => void) | null = null;
  onerror: ((event: Event) => void) | null = null;

  send(data: string): void {
    console.log('MockWebSocket send:', data);
    // 模拟消息回复
    setTimeout(() => {
      if (this.onmessage) {
        const parsedData = JSON.parse(data);
        if (parsedData.channel === 'chat' && parsedData.type === 'text') {
          // 模拟AI回复
          this.onmessage(new MessageEvent('message', {
            data: JSON.stringify({
              channel: 'chat',
              type: 'text',
              clientId: 'mock-client',
              body: '这是一个模拟的回复消息。'
            })
          }));
          
          // 模拟消息结束
          setTimeout(() => {
            this.onmessage?.(new MessageEvent('message', {
              data: JSON.stringify({
                channel: 'chat',
                type: 'end',
                clientId: 'mock-client',
                body: ''
              })
            }));
          }, 500);
        }
      }
    }, 300);
  }

  close(): void {
    console.log('MockWebSocket closed');
    if (this.onclose) {
      this.onclose(new CloseEvent('close'));
    }
  }
}

const useStore = create<SharedStore>((set, get) => ({
  // 聊天设置
  chatListStyle: 'default',
  chatStream: true,
  setChatListStyle: (style) => {
    if (typeof window !== 'undefined') {
      localStorage.setItem('chatListStyle', style);
    }
    set({ chatListStyle: style });
  },
  setChatStream: (stream) => {
    if (typeof window !== 'undefined') {
      localStorage.setItem('chatStream', String(stream));
    }
    set({ chatStream: stream });
  },
  
  // 登录相关
  showLoginDialog: false,
  setShowLoginDialog: (show) => set({ showLoginDialog: show }),
  
  // WebSocket 相关
  socket: {
    conn: null,
    isConnected: false,
    messageHandlers: new Map(),
  },
  connectWebSocket: () => {
    const { socket } = get();
    if (socket.conn) {
      socket.conn.close();
    }
    
    // 使用模拟的 WebSocket
    const conn = new MockWebSocket() as unknown as WebSocket;
    
    conn.onopen = () => {
      set({
        socket: {
          ...socket,
          conn,
          isConnected: true,
        }
      });
    };
    
    conn.onclose = () => {
      set({
        socket: {
          ...socket,
          isConnected: false,
        }
      });
    };
    
    conn.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data);
        const channel = data.channel;
        const handlers = socket.messageHandlers.get(channel) || [];
        handlers.forEach(handler => handler(data));
      } catch (e) {
        console.error('Failed to parse WebSocket message:', e);
      }
    };
    
    set({
      socket: {
        ...socket,
        conn,
      }
    });
  },
  disconnectWebSocket: () => {
    const { socket } = get();
    if (socket.conn) {
      socket.conn.close();
    }
    set({
      socket: {
        ...socket,
        conn: null,
        isConnected: false,
      }
    });
  },
  addMessageHandler: (channel, handler) => {
    const { socket } = get();
    const handlers = socket.messageHandlers.get(channel) || [];
    handlers.push(handler);
    socket.messageHandlers.set(channel, handlers);
    set({ socket });
  },
  removeMessageHandler: (channel) => {
    const { socket } = get();
    socket.messageHandlers.delete(channel);
    set({ socket });
  }
}));

export const useSharedStore = (selector: (state: SharedStore) => any) => {
    const state = useStore(selector);
    const [hydrated, setHydrated] = React.useState(false);

    React.useEffect(() => {
        // 从 localStorage 初始化状态，确保只在客户端执行一次
        useStore.setState({
            chatListStyle: localStorage.getItem('chatListStyle') || 'default',
            chatStream: localStorage.getItem('chatStream') !== 'false'
        });
        setHydrated(true);
    }, []);

    const initialState: SharedStore = {
        chatListStyle: 'default',
        chatStream: true,
        setChatListStyle: () => {},
        setChatStream: () => {},
        showLoginDialog: false,
        setShowLoginDialog: () => {},
        socket: {
            conn: null,
            isConnected: false,
            messageHandlers: new Map(),
        },
        connectWebSocket: () => {},
        disconnectWebSocket: () => {},
        addMessageHandler: () => {},
        removeMessageHandler: () => {},
    };

    if (!hydrated) {
        return selector(initialState);
    }

    return state;
}; 